/**
 * Copyright 2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "kernel_builder/aicpu_ascend_ops_kernel_builder.h"
#include <memory>
#include <vector>
#include "common/config/config_file.h"
#include "common/util/error_manager/error_manager.h"
#include "common/util/util.h"
#include "error_code/error_code.h"
#include "ge/ge_api_types.h"
#include "register/ops_kernel_builder_registry.h"

namespace aicpu {
REGISTER_OPS_KERNEL_BUILDER(kAicpuOpsKernelInfo, AicpuAscendOpsKernelBuilder);
ge::Status AicpuAscendOpsKernelBuilder::Initialize(
    const std::map<std::string, std::string> &options) {
  AICPUE_LOGI("Begin to initialize aicpu ops kernel builder")

  AICPU_CHECK_RES(Finalize())

  std::string kernel_builder_str = "CUSTAICPUBuilder,AICPUBuilder";
  std::vector<std::string> kernel_builders;
  ConfigFile::GetInstance().SplitValue(kernel_builder_str, kernel_builders);
  // must be not empty...
  if (kernel_builders.empty()) {
    AICPUE_LOGE("KernelBuilders is empty.");
    return NONE_KERNEL_BUILDER;
  }
  for (auto kernel_builder : kernel_builders) {
    FACTORY_KERNEL_BUILDER::FactoryType kernel_builder_ptr =
        FACTORY_KERNEL_BUILDER::Produce(kernel_builder);
    if (kernel_builder_ptr == nullptr) {
      AICPU_REPORT_INNER_ERROR("Crate %s failed.", kernel_builder.c_str());
      return KERNEL_BUILDER_INSTANCE_FAILED;
    }
    kernel_builder_map_[kernel_builder] = kernel_builder_ptr;
    AICPU_CHECK_RES(kernel_builder_ptr->Initialize());
  }
  return ge::SUCCESS;
}

ge::Status AicpuAscendOpsKernelBuilder::Finalize() {
  kernel_builder_map_.clear();
  return ge::SUCCESS;
}

ge::Status AicpuAscendOpsKernelBuilder::CalcOpRunningParam(ge::Node &node) {
  ge::OpDescPtr op_desc_ptr = node.GetOpDesc();
  AICPU_CHECK_NOTNULL_ERRCODE(op_desc_ptr, ErrorCode::INPUT_PARAM_NULL)
  std::string op_type = op_desc_ptr->GetType();
  if (op_type == kFrameworkOp) {
    std::string original_type;
    AICPU_IF_BOOL_EXEC(
        (!ge::AttrUtils::GetStr(op_desc_ptr, kOriginalType, original_type)),
        AICPU_REPORT_CALL_ERROR("Get attr[%s] of op[%s] failed.",
            kOriginalType.c_str(), node.GetName().c_str());
        return ErrorCode::GET_ORIGINAL_TYPE_FAILED)
    if (original_type.empty()) {
      AICPU_REPORT_INNER_ERROR("Attr[%s] of op[%s] is empty.",
                kOriginalType.c_str(), node.GetName().c_str());
      return STR_IS_EMPTY;
    }
    op_desc_ptr->SetType(original_type);
    op_type = original_type;
  }
  std::string kernel_name;
  (void)ge::AttrUtils::GetStr(op_desc_ptr, "opKernelLib", kernel_name);
  const KernelBuilderPtr &kernel_builder = GetKernelBuilderByName(kernel_name);

  AICPU_CHECK_NOTNULL_ERRCODE(kernel_builder, ErrorCode::NONE_KERNEL_BUILDER)
  return kernel_builder->CalcOpRunningParam(node);
}

ge::Status AicpuAscendOpsKernelBuilder::GenerateTask(
    const ge::Node &node, ge::RunContext &context,
    std::vector<domi::TaskDef> &tasks) {
  ge::OpDescPtr op_desc_ptr = node.GetOpDesc();
  AICPU_CHECK_NOTNULL_ERRCODE(op_desc_ptr, ErrorCode::INPUT_PARAM_NULL);

  // get original type
  std::string op_type = op_desc_ptr->GetType();
  if (op_type == kFrameworkOp) {
    std::string original_type;
    AICPU_IF_BOOL_EXEC(
        (!ge::AttrUtils::GetStr(op_desc_ptr, kOriginalType, original_type)),
         AICPU_REPORT_CALL_ERROR("Get attr[%s] of op[%s] failed.",
            kOriginalType.c_str(), op_desc_ptr->GetName().c_str());
        return ErrorCode::GET_ORIGINAL_TYPE_FAILED)
    op_type = original_type;
  }

  std::string kernel_name;
  (void)ge::AttrUtils::GetStr(op_desc_ptr, "opKernelLib", kernel_name);
  const KernelBuilderPtr &kernel_builder = GetKernelBuilderByName(kernel_name);
  AICPU_CHECK_NOTNULL_ERRCODE(kernel_builder, ErrorCode::NONE_KERNEL_BUILDER);
  return kernel_builder->GenerateTask(node, context, tasks);
}

KernelBuilderPtr AicpuAscendOpsKernelBuilder::GetKernelBuilderByName(
    const std::string &kernel_name) {
  unsigned int index = kernel_name.find("Kernel");
  if (index == kernel_name.npos) {
    AICPU_REPORT_INNER_ERROR("Wrong kernelName[%s].", kernel_name.c_str());
  }
  std::string kernel_builder_name = kernel_name.substr(0, index) + "Builder";
  if (kernel_builder_map_.find(kernel_builder_name) !=
      kernel_builder_map_.end()) {
    return kernel_builder_map_[kernel_builder_name];
  }
   AICPU_REPORT_INNER_ERROR("%s kernel builder is not exist.",
      kernel_builder_name.c_str());
  return nullptr;
}
}  // namespace aicpu
